#!/usr/bin/env python
from Bio import AlignIO
#from Bio import SeqIO
#from Bio.PDB import *
#from Bio.SeqRecord import SeqRecord
#from Bio.Seq import Seq
#import os
import re
import sys
import pickle
from collections import defaultdict
#from bio import blosum
from itertools import combinations
from pprint import pprint
def is_contact(cont, pdbid, peppos, dist, id, residue):
    #flag = False
    GLGF = 22
    GD = 55
    for j in cont[peppos]:
        if j[0] == "GLGF":
            num = id - GLGF +1 
        elif j[0] == "GD":
            num = id - GD + 1
        
        if j[1] == num:
            #print "num:", num
            for i in cont[peppos][j]:
                #print i
                if i[0] == pdbid:
                    if residue == i[1]:
                        if i[2] <= dist:
                            return True
                    else:
                        print "residues don't match.", pdbid, j, i[1], id, residue
    return False

def pretty_aln(aln, cont, dist):
    #printing the alignment
    rich_alignment = dict()
    #format should be
    #rich_alignment[pdbid] = [[res1, [contating_peptide positions1, 2...]],[res2, [cp1, 2 ...]], []...]
    gaps_count = dict()
    for i in range(aln.get_alignment_length()):
        
        for ent in aln:
            residue = ent.seq[i]
            contact_pep = []
            
            pdbid = ent.id
            
            if gaps_count.has_key(pdbid):
                pointer =  i - gaps_count[pdbid]
            else:
                gaps_count[pdbid] = 0
             
            if ent.seq[i] == '-' or ent.seq[i] =='.':
                gaps_count[pdbid] += 1
                #print pdbid, gaps_count[pdbid], i, ent.seq[i]
            else:
                #ms = missres[pdbid]
                #print pdbid, pointer, ent.seq[i], tpos, tres
                for peppos in range(4):
                    #print pdbid
                    if is_contact(cont, pdbid, peppos, dist, i, residue):
                        
                        contact_pep.append(peppos)
                        #print contact_pep
            if rich_alignment.has_key(pdbid):
                rich_alignment[pdbid].append([residue, contact_pep])
            else:
                rich_alignment[pdbid] = [[residue, contact_pep]]
            
    return rich_alignment


def print_aln(aln, peppos):
    CSI = "\x1B["
    csi_reset = CSI + "0m"
    aln_len = 0
    for key in aln.keys():
        aln_len = max(len(aln[key]), aln_len)
    respline = 80
    lastline = aln_len % respline
    iter_range2 = range(0, aln_len, respline)
    iter_range = iter_range2
    iter_range2.append(iter_range2[-1]+lastline)
    for i in range(len(iter_range)-1):
        #print i
        for key in aln.keys():
            #print iter_range[i], iter_range2[i+1]
            sys.stdout.write(key+" ")
            for p in range(iter_range[i], iter_range2[i+1]):
                #print p
                #peppos = -1  print all contact residues
                if peppos == -1 and aln[key][p][1] != []:
                    sys.stdout.write(CSI+"41m"+aln[key][p][0]+csi_reset)   
                elif peppos in aln[key][p][1]:
                    sys.stdout.write(CSI+"41m"+aln[key][p][0]+csi_reset)
                else:
                    sys.stdout.write(aln[key][p][0])
            print
        print


#load files
contactfile = "contactinfo.dat"
alnfile = "tmp_aln.fasta"

with open(contactfile, 'rb') as pkl_file: 
    contactdat, pos_glgf, pos_gd = pickle.load(pkl_file)
aln = AlignIO.read(alnfile, "fasta")
    
rich_aln = pretty_aln(aln, contactdat, 6)

print_aln(rich_aln, -1)
#pprint( contactdat[0])
'''
final_out = {
    0: {(GLGF,1): [[PDBID],[RES],[DIST]]
        (GLGF,2): [[],[],[]]},
    1: {},
    2: {},
    3: {},
    4: {},
}
'''
    